import math
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped, Quaternion
from std_msgs.msg import Float64MultiArray
from nav_msgs.msg import Path, Odometry
from rclpy.publisher import Publisher
import numpy as np
from rclpy.qos import DurabilityPolicy, QoSProfile
from message_filters import Subscriber, TimeSynchronizer
import tf_transformations as tft


class LocalPathPlaner(Node):
    """
    Local Path Planner
    step1 : generate path from start to end
        functin1: similar_triangle
            can let middle stage of path be straight line, which is necessary for dock together
        function2: bezier_curve
    step2 : plan pitch and angle
        function1: trapezoidal planning
        functionn2: mpc
    """

    def __init__(self):
        super().__init__("local_path_planner")
        startSub = Subscriber(self, Odometry, "odom_start")
        endSub = Subscriber(self, Odometry, "odom_end")
        self.odomTs = TimeSynchronizer([startSub, endSub], 10)
        self.odomTs.registerCallback(self.odom_callback)
        qos = QoSProfile(depth=1, durability=DurabilityPolicy.TRANSIENT_LOCAL)
        self.path_pub: Publisher = self.create_publisher(Path, "local_path", qos)
        self.vec_pub: Publisher = self.create_publisher(
            Float64MultiArray, "local_vec", qos
        )
        self.declare_parameter("radius", 0.2)
        self.radius: float = self.get_parameter("radius").value
        self.declare_parameter("invert", False)
        self.invert: bool = self.get_parameter("invert").value
        self.declare_parameter("step_length", 0.1)
        self.stepLength: float = self.get_parameter("step_length").value
        self.declare_parameter("start_vec", 1.0)
        self.declare_parameter("dock_vec", 0.1)
        self.startVec: float = self.get_parameter("start_vec").value
        self.dockVec: float = self.get_parameter("dock_vec").value

        # for test
        # self.test_env()

    def odom_callback(self, start: Odometry, end: Odometry):
        # self.get_logger().info("odom changed, replan path")
        path, vec = self.solution1(
            start,
            end,
            self.radius,
            -self.startVec if self.invert else self.startVec,
            -self.dockVec if self.invert else self.dockVec,
        )
        self.path_pub.publish(path)
        self.vec_pub.publish(vec)

    def solution1(
        self,
        start: Odometry,
        end: Odometry,
        l: float,
        start_vel: float,
        dock_vec: float,
    ) -> tuple[Path, Float64MultiArray]:
        """
        起点和终点以及初始朝向的交点构成三角形，实际轨迹为与三角形相似，但形状更小
        """
        startP = np.array([start.pose.pose.position.x, start.pose.pose.position.y])
        endP = np.array([end.pose.pose.position.x, end.pose.pose.position.y])

        distance = np.linalg.norm(endP - startP)
        l = distance * l

        startQua = (
            start.pose.pose.orientation.x,
            start.pose.pose.orientation.y,
            start.pose.pose.orientation.z,
            start.pose.pose.orientation.w,
        )
        startDir = tft.quaternion_matrix(startQua)[0:2, 0]
        startDir /= np.linalg.norm(startDir)
        endQua = (
            end.pose.pose.orientation.x,
            end.pose.pose.orientation.y,
            end.pose.pose.orientation.z,
            end.pose.pose.orientation.w,
        )
        endDir = tft.quaternion_matrix(endQua)[0:2, 0]
        endDir /= np.linalg.norm(endDir)

        if False and np.arccos(np.dot(startDir, endDir)) < np.radians(5):
            points = self.besizer_curve(startP, (startP + endP) / 2, endP, 10)
        else:
            if self.invert:
                startMid = startP - startDir * l
                endMid = endP + endDir * l
                midDir = endMid - startMid
                midDir /= np.linalg.norm(midDir)
                startMidInLine = startMid + midDir * l
                endMidInLine = endMid - midDir * l
                mid = (startMidInLine + endMidInLine) / 2
            else:
                startMid = startP + startDir * l
                endMid = endP - endDir * l
                midDir = endMid - startMid
                midDir /= np.linalg.norm(midDir)
                startMidInLine = startMid + midDir * l
                endMidInLine = endMid - midDir * l
                mid = (startMidInLine + endMidInLine) / 2

            startPoints = self.besizer_curve(
                startP,
                startMid,
                startMidInLine,
                int(np.linalg.norm(startP - startMidInLine) / self.stepLength),
            )
            endPoints = self.besizer_curve(
                endMidInLine,
                endMid,
                endP,
                int(np.linalg.norm(endMidInLine - endP) / self.stepLength),
            )
            midPoints = self.besizer_curve(
                startMidInLine,
                mid,
                endMidInLine,
                int(np.linalg.norm(startMidInLine - endMidInLine) / self.stepLength),
            )
            points = np.vstack([startPoints, midPoints[1:-1, :], endPoints])

        # points = [startP, startMid, startMidInLine, mid, endMidInLine, endMid, endP]
        velocity = np.linspace(start_vel, dock_vec, len(points))

        path = Path()
        path.header.frame_id = "map"
        path.header.stamp = self.get_clock().now().to_msg()
        for i in range(len(points)):
            pose = PoseStamped()
            pose.header.frame_id = "map"
            pose.header.stamp = self.get_clock().now().to_msg()
            pose.pose.position.x = points[i][0]
            pose.pose.position.y = points[i][1]
            pose.pose.position.z = 0.0
            if i == len(points) - 1:
                pose.pose.orientation = end.pose.pose.orientation
            else:
                dx = points[i + 1][0] - points[i][0]
                dy = points[i + 1][1] - points[i][1]
                # 插值计算点的朝向
                qua = tft.quaternion_from_euler(
                    0.0,
                    0.0,
                    (
                        math.atan2(
                            -dy,
                            -dx,
                        )
                        if self.invert
                        else math.atan2(dy, dx)
                    ),
                )
                pose.pose.orientation = Quaternion(
                    x=qua[0], y=qua[1], z=qua[2], w=qua[3]
                )
            path.poses.append(pose)
        vec = Float64MultiArray()
        vec.data = velocity
        return path, vec

    def besizer_curve(self, p0, p1, p2, num_points):
        t = np.linspace(0, 1, num_points)[:, np.newaxis]
        points = (1 - t) ** 2 * p0 + 2 * t * (1 - t) * p1 + t**2 * p2
        return points

    def test_env(self):
        start = Odometry()
        start.header.frame_id = "map"
        start.header.stamp = self.get_clock().now().to_msg()
        start.pose.pose.position.x = 0.0
        start.pose.pose.position.y = 0.0
        qua = tft.quaternion_from_euler(0, 0, 0)
        start.pose.pose.orientation = Quaternion(x=qua[0], y=qua[1], z=qua[2], w=qua[3])
        end = Odometry()
        end.header.frame_id = "map"
        end.header.stamp = self.get_clock().now().to_msg()
        end.pose.pose.position.x = 10.0
        end.pose.pose.position.y = 10.0
        qua = tft.quaternion_from_euler(0, 0, math.pi / 2)
        end.pose.pose.orientation = Quaternion(x=qua[0], y=qua[1], z=qua[2], w=qua[3])
        self.path_pub.publish(self.solution1(start, end, 1))


def main(args=None):
    rclpy.init(args=args)
    local_path_planer = LocalPathPlaner()
    rclpy.spin(local_path_planer)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
